#pragma OPENCL EXTENSION cl_khr_fp16 : enable
#pragma OPENCL EXTENSION cl_khr_subgroups : enable

kernel void kernel_soft_max(
        global char * src0,
        ulong offset0,
        global char * src1,
        ulong offset1,
        global char * src2,
        ulong offset2,
        global char * dst,
        ulong offsetd,
        int ne00,
        ulong nb01,
        ulong nb02,
        ulong nb03,
        int ne12,
        int ne13,
        ulong nb11,
        ulong nb12,
        ulong nb13,
        ulong nb1,
        ulong nb2,
        ulong nb3,
        float scale,
        float max_bias,
        float m0,
        float m1,
        int n_head_log2
) {
    src0 = src0 + offset0;
    src1 = src1 + offset1;
    src2 = src2 + offset2;
    dst  = dst  + offsetd;

    int i03 = get_group_id(2);
    int i02 = get_group_id(1);
    int i01 = get_group_id(0);

    int i13 = i03%ne13;
    int i12 = i02%ne12;
    int i11 = i01;

    // ALiBi
    float slope = 1.0f;
    if (max_bias > 0.0f) {
        int h = i02;
        float base = h < n_head_log2 ? m0 : m1;
        int   exp  = h < n_head_log2 ? h + 1 : 2*(h - n_head_log2) + 1;
        slope = pow(base, exp);
    }

    global float * psrc0 = (global float *)(src0 + i01*nb01 + i02*nb02 + i03*nb03);
    global float * pdst  = (global float *)(dst  + i01*nb1 + i02*nb2 + i03*nb3);
    global float * psrc2;
    global float * pmask;

    float lmax = 0.0f;
    float max = 0.0f;
    float lsum = 0.0f;
    float sum = 0.0f;
    float exp_psrc0 = 0.0f;
    float p_scale=0.0f;
    float s_pmask=0.0f;

    if (src2 != src0) {
        // initial max
        psrc2 = (global float *)src2;
        if (psrc2) {
            lmax = psrc2[i02];
        }
        else{
            lmax = -INFINITY;
        }

        // parallel max
        if (src1 != src0) {
            pmask = (global float *)(src1 + i11*nb11 + i12*nb12 + i13*nb13);
            printf("[softmax complex] pmask : %f \n", *pmask);

            // first pass : compute max
            for (int i00 = get_local_id(0); i00 < ne00; i00 += get_local_size(0)) {
                p_scale = psrc0[i00]*scale;
                s_pmask = slope*pmask[i00];

                if (pmask) {
                    lmax = fmax(lmax, p_scale + s_pmask);
                }else{
                    lmax = fmax(lmax, p_scale + 0.0f);
                }
            }
            max = sub_group_reduce_max(lmax);

            // second pass : compute sum
            for (int i00 = get_local_id(0); i00 < ne00; i00 += get_local_size(0)) {
                p_scale = psrc0[i00]*scale;
                s_pmask = slope*pmask[i00];

                if (pmask) {
                    exp_psrc0 = exp(p_scale + s_pmask - max);
                    lsum += exp_psrc0;
                    pdst[i00] = exp_psrc0;
                }else{
                    exp_psrc0 = exp(p_scale + 0.0f - max);
                    lsum += exp_psrc0;
                    pdst[i00] = exp_psrc0;
                }
            }
            sum = sub_group_reduce_add(lsum);
            if (psrc2) {
                sum += exp(psrc2[i02] - max);
            }

            // third pass : normalize
            for (int i00 = get_local_id(0); i00 < ne00; i00 += get_local_size(0)) {
                pdst[i00] /= sum;
            }
            
        } else {
            pmask = (global float *)0;   // 避免隐式 addrspace cast
            printf("[softmax complex] pmask : %f \n", *pmask);

            // first pass : compute max
            for (int i00 = get_local_id(0); i00 < ne00; i00 += get_local_size(0)) {
                p_scale = psrc0[i00]*scale;
                s_pmask = slope*pmask[i00];

                if (pmask) {
                    lmax = fmax(lmax, p_scale + s_pmask);
                }else{
                    lmax = fmax(lmax, p_scale + 0.0f);
                }
            }
            max = sub_group_reduce_max(lmax);

            // second pass : compute sum
            for (int i00 = get_local_id(0); i00 < ne00; i00 += get_local_size(0)) {
                p_scale = psrc0[i00]*scale;
                s_pmask = slope*pmask[i00];

                if (pmask) {
                    exp_psrc0 = exp(p_scale + s_pmask - max);
                    lsum += exp_psrc0;
                    pdst[i00] = exp_psrc0;
                }else{
                    exp_psrc0 = exp(p_scale + 0.0f - max);
                    lsum += exp_psrc0;
                    pdst[i00] = exp_psrc0;
                }
            }
            sum = sub_group_reduce_add(lsum);
            if (psrc2) {
                sum += exp(psrc2[i02] - max);
            }

            // third pass : normalize
            for (int i00 = get_local_id(0); i00 < ne00; i00 += get_local_size(0)) {
                pdst[i00] /= sum;
            }

        }
   
    } else {
        psrc2 = (global float *)0;   // 避免隐式 addrspace cast

        // initial max
        if (psrc2) {
            lmax = psrc2[i02];
        } else {
            lmax = -INFINITY;
        }

        // parallel max
        if (src1 != src0) {
            pmask = (global float *)(src1 + i11*nb11 + i12*nb12 + i13*nb13);
            printf("[softmax complex] pmask : %f \n", *pmask);

            // first pass : compute max
            for (int i00 = get_local_id(0); i00 < ne00; i00 += get_local_size(0)) {
                p_scale = psrc0[i00]*scale;
                s_pmask = slope*pmask[i00];

                if (pmask) {
                    lmax = fmax(lmax, p_scale + s_pmask);
                }else{
                    lmax = fmax(lmax, p_scale + 0.0f);
                }
            }
            max = sub_group_reduce_max(lmax);

            // second pass : compute sum
            for (int i00 = get_local_id(0); i00 < ne00; i00 += get_local_size(0)) {
                p_scale = psrc0[i00]*scale;
                s_pmask = slope*pmask[i00];

                if (pmask) {
                    exp_psrc0 = exp(p_scale + s_pmask - max);
                    lsum += exp_psrc0;
                    pdst[i00] = exp_psrc0;
                }else{
                    exp_psrc0 = exp(p_scale + 0.0f - max);
                    lsum += exp_psrc0;
                    pdst[i00] = exp_psrc0;
                }
            }
            sum = sub_group_reduce_add(lsum);
            if (psrc2) {
                sum += exp(psrc2[i02] - max);
            }

            // third pass : normalize
            for (int i00 = get_local_id(0); i00 < ne00; i00 += get_local_size(0)) {
                pdst[i00] /= sum;
            }
            
        } else {
            pmask = (global float *)0;   // 避免隐式 addrspace cast
            printf("[softmax complex] pmask : %f \n", *pmask);

            // first pass : compute max
            for (int i00 = get_local_id(0); i00 < ne00; i00 += get_local_size(0)) {
                p_scale = psrc0[i00]*scale;
                s_pmask = slope*pmask[i00];

                if (pmask) {
                    lmax = fmax(lmax, p_scale + s_pmask);
                }else{
                    lmax = fmax(lmax, p_scale + 0.0f);
                }
            }
            max = sub_group_reduce_max(lmax);

            // second pass : compute sum
            for (int i00 = get_local_id(0); i00 < ne00; i00 += get_local_size(0)) {
                p_scale = psrc0[i00]*scale;
                s_pmask = slope*pmask[i00];

                if (pmask) {
                    exp_psrc0 = exp(p_scale + s_pmask - max);
                    lsum += exp_psrc0;
                    pdst[i00] = exp_psrc0;
                }else{
                    exp_psrc0 = exp(p_scale + 0.0f - max);
                    lsum += exp_psrc0;
                    pdst[i00] = exp_psrc0;
                }
            }
            sum = sub_group_reduce_add(lsum);
            if (psrc2) {
                sum += exp(psrc2[i02] - max);
            }

            // third pass : normalize
            for (int i00 = get_local_id(0); i00 < ne00; i00 += get_local_size(0)) {
                pdst[i00] /= sum;
            }

        }
    }

}